# -*- coding: utf-8 -*-
# FileName:     scope.py
# time:         23/4/24 024 下午 7:13
# Author:       Zhou Hang
# Description:  30-直角三角形对称，旋转，阵列化生成纹理图形

import numpy as np
import cv2
import math
from pathlib import *

from core.rec_matrix import rec_matrix


def make_border_lr(src, left, right):
    """
    在src的左右填充黑边来扩展大小
    :param src: 待扩展的图片
    :param left: 左边扩展的长度
    :param right: 右边扩展的长度
    :return: 扩展后的图形
    """
    return cv2.copyMakeBorder(src, 0, 0, left, right, cv2.BORDER_CONSTANT, value=[0, 0, 0])


source_file = Path("./src.jpg")  # 源文件
if not source_file.exists():
    print("原始图像不存在，请将待处理文件重命名为src.jpg, 放在当前目录下")
    exit(0)

temp_path = Path("./resources/temp")  # 中间结果目录
result_path = Path("./resources/res")  # 结果存放目录
temp_path.mkdir(parents=True, exist_ok=True)
result_path.mkdir(parents=True, exist_ok=True)

# ================================= 截取出指定像素 =================================
img = cv2.imread(str(source_file))
print(f"源文件横坐标x范围: 0-{img.shape[1]} 纵坐标y范围0-{img.shape[0]}")
# 右上角的三角形
lt_x = int(input("请输入直角三角形左上角横坐标x: "))
lt_y = int(input("请输入直角三角形左上角纵坐标y: "))
target_width = int(input("请输入截取三角形最短边长度: "))
# lt_x, lt_y = 500, 1000  # 待切割矩形的左上角坐标(三角形长直角边对应的点)
# target_width = 100  # 三角形最短边的边长

target_height = math.ceil(target_width * math.sqrt(3))
uniq_name = f"x{lt_x}_y{lt_y}_l{target_width}"
final_path = result_path / f"{uniq_name}.jpg"  # 最终结果存放路径

triangle = np.array([
    [lt_x, lt_y],
    [lt_x + target_width, lt_y],
    [lt_x + target_width, lt_y + target_height]],
    dtype=np.int32)
triangle_mask = np.zeros(img.shape, dtype="uint8")
_ = cv2.fillPoly(triangle_mask, [triangle], [255, 255, 255])
img = cv2.bitwise_and(img, triangle_mask)
sample_tri = img[lt_y: lt_y + target_height, lt_x: lt_x + target_width]

min_path = temp_path / f"{uniq_name}_min.jpg"
try:
    cv2.imwrite(str(min_path), sample_tri)
except cv2.error:
    print("截取的三角形超过原始图像范围，请修改")
    exit(1)
cont = input(f"截取图像已经保存至 {min_path}，请查看，并选择是否继续？（Y/N）")
if not (cont == "Y" or cont == "y"):
    exit(0)

sample_tri_h = cv2.flip(sample_tri, 1)  # 水平翻转

# ================================= 生成所有部分 =================================
tri1 = cv2.hconcat([sample_tri, sample_tri_h])
height, width = tri1.shape[:2]

side_length = width
half_side = side_length // 2

# 将图像增加黑边扩充大小，一遍获得旋转后的图像
tri1 = make_border_lr(tri1, width // 2, width // 2)

_, width = tri1.shape[:2]
x0, y0 = width // 2, height

M60 = cv2.getRotationMatrix2D((x0, y0), 60, 1.0)
tri42 = cv2.warpAffine(tri1, M60, (width, height))

M60_n = cv2.getRotationMatrix2D((x0, y0), -60, 1.0)
tri35 = cv2.warpAffine(tri1, M60_n, (width, height))

x0, y0 = width // 2, height // 2
M180 = cv2.getRotationMatrix2D((x0, y0), 180, 1.0)

tri11 = cv2.warpAffine(tri1, M180, (width, height))
tri2244 = cv2.warpAffine(tri42, M180, (width, height))
tri5533 = cv2.warpAffine(tri35, M180, (width, height))
tri33 = tri5533[:, side_length // 2: side_length]
tri55 = tri5533[:, 0: side_length // 2]
tri2 = tri42[:, side_length // 2: side_length]
tri4 = tri42[:, 0: side_length // 2]

# 将所有部分填充到同样大小
tri1t = make_border_lr(tri1, side_length, 0)
tri1b = make_border_lr(tri1[:, half_side:], 0, side_length + half_side)
tri42t = make_border_lr(tri42, side_length, 0)
tri35t = make_border_lr(tri35, side_length, 0)
tri35b = make_border_lr(tri35[:, half_side:], 0, side_length + half_side)
tri2244t = make_border_lr(tri2244[:, half_side:], 0, side_length + half_side)
tri2244b = make_border_lr(tri2244, side_length, 0)
tri11t = make_border_lr(tri11[:, half_side:], 0, side_length + half_side)
tri11b = make_border_lr(tri11, side_length, 0)
tri33t = make_border_lr(tri33, 0, 2 * side_length + half_side)
tri55t = make_border_lr(tri55, 2 * side_length + half_side, 0)
tri2b = make_border_lr(tri2, 0, 2 * side_length + half_side)
tri4b = make_border_lr(tri4, 2 * side_length + half_side, 0)
tri5533b = make_border_lr(tri5533, side_length, 0)

# ================================= 将所有部分拼接到一起 =================================
height, width, _ = tri1t.shape

t = tri1t + tri42t + tri35t + tri2244t + tri33t + tri55t + tri11t
b = tri1b + tri35b + tri2244b + tri11b + tri2b + tri4b + tri5533b
# 高斯滤波，中值滤波，方框滤波，经过测试中值滤波最为有效去除黑边
t = cv2.medianBlur(t, 5)
b = cv2.medianBlur(b, 5)

unit = cv2.vconcat([t, b])  # 将上下半部分拼接起来，生成最终的单元矩阵

# ================================= 膨胀消去黑点 =================================
kernel = np.ones((2, 2), np.uint8)
unit = cv2.dilate(unit, kernel, iterations=2)
unit_path = temp_path / f"{uniq_name}_unit.jpg"
cv2.imwrite(str(unit_path), unit)

# ================================= 对生成的矩形阵列拼接 =================================
rec = unit  # just 换一个名字
new_rec = rec_matrix(rec, 5, 3)
# 某些情况下阵列完后存在黑边，再膨胀一次
new_rec = cv2.dilate(new_rec, kernel, iterations=2)
cv2.imwrite(str(final_path), new_rec)
print(f"生成完毕，图片保存至: {final_path}")
